import pandas as pd
import numpy as np
import h5py
from utils.config import *
from utils.data_loader import day_block_shuffle, process_images
from utils.train import train_model
from utils.evaluate import evaluate_model, plot_results, plot_sunny_results, plot_cloudy_results
from learner.measure import *
from learner.train import get_data, train_model_and_predict, determine_final_prediction, save_classified_output, dates_sort
from processing.forecast_classification import prepare_data

   
#其他参数配置请查看utils/config文件
#可选择参数:["MLP", "LSTM", "CNN_LSTM"]
config.MODEL_SELECT = "LSTM"

if __name__ == "__main__":
    """
    1.深度森林算法->光伏分类(PV-classification)->(0:晴天/1:阴天)
    """
    data_name = 'solar_data'
    #获取训练数据和测试数据
    train_dates, train_data, train_label, test_dates, test_data, test_label = get_data(data_name)
    #深度森林训练轮数
    epoch = 5
    #模型的训练与预测，预测五轮的结果
    all_predictions = train_model_and_predict(data_name, train_data, train_label, test_data, "one-error", epoch)
    #对五轮的预测结果取众数
    prediction = determine_final_prediction(all_predictions)
    print("true:", test_label)
    print("prediction:", prediction)
    #保存预测结果
    save_classified_output(train_dates, train_label, test_dates, test_label, prediction, data_name)
    #将预测结果按日期排序并分为晴天和阴天
    prediction = dates_sort(test_label, prediction, test_dates)
    for index, pred_value in enumerate(prediction):
        if pred_value == 0:
            config.SUNNY_DATES.append(config.TEST_DATES[index])
        else:
            config.CLOUDY_DATES.append(config.TEST_DATES[index])
    #根据深度森林预测的结果准备深度学习模型训练和测试所需数据
    if not all([os.path.exists(path) for path in config.file_paths]):
        prepare_data()
    
    """ 
    2.深度学习模型->光伏预测(PV-forcasting) 
    """
    # 配置GPU，没有GPU将用CPU训练
    configure_hardware()
    # 三种情况下训练三个模型（阴天 晴天 阴天+晴天）
    weathers = ["CLOUDY", "SUNNY", "ALL"]
    for weather_select in weathers:
        # 配置天气模式
        config.WEATHER_SELECT = weather_select
        # 读取数据集训练
        data_path = os.path.join(config.data_folder, "forecast_dataset.hdf5")
        times_trainval = np.load(os.path.join(config.data_folder, "times_trainval.npy"), allow_pickle=True)
        times_test = np.load(os.path.join(config.data_folder, "times_test.npy"), allow_pickle=True)
        
        # 训练模型
        history, model = train_model(data_path, times_trainval)

        # 读取测试集
        with h5py.File(data_path, 'r') as f:
            test_images = process_images(f['test']['image_log_test'][...])
            test_pv_log = f['test']['pv_log_test'][...].astype('float32')
            test_pv_pred = f['test']['pv_pred_test'][...].astype('float32')
        # 模型测试与评估
        metrics, predictions = evaluate_model(model, (test_images, test_pv_log, test_pv_pred), times_test)

        sunny_dates = [datetime.date(*d) for d in config.SUNNY_DATES]
        cloudy_dates = [datetime.date(*d) for d in config.CLOUDY_DATES]
        
        # 只评估RMSE，若需其他评价指标请自行添加
        # 模型输出图片位于目录: solar/data/model_output/SUNSET_forecast_2017_2019_data
        if weather_select == "ALL":
            print(f"{weather_select} Test RMSE: {metrics['overall']['RMSE']:.3f}")
            print(f"{weather_select} Sunny Days RMSE: {metrics['sunny']['RMSE']:.3f}")
            print(f"{weather_select} Cloudy Days RMSE: {metrics['cloudy']['RMSE']:.3f}")
            plot_results(times_test, test_pv_pred, predictions, sunny_dates, cloudy_dates)
        elif weather_select == "SUNNY":
            print(f"{weather_select} Sunny Days RMSE: {metrics['sunny']['RMSE']:.3f}")
            plot_sunny_results(times_test, test_pv_pred, predictions, sunny_dates)
        else:
            print(f"{weather_select} Cloudy Days RMSE: {metrics['cloudy']['RMSE']:.3f}")
            plot_cloudy_results(times_test, test_pv_pred, predictions, cloudy_dates)
            
            
            